#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <sys/time.h>
#include <upc.h>

#ifndef NUM_THREADS
#define NUM_THREADS     2
#endif
#ifndef CAPACITY
#define CAPACITY        999
#endif
#ifndef NITEMS
#define NITEMS          5000
#endif
#define BS_ITEMS        ((NITEMS+NUM_THREADS-1)/NUM_THREADS)
#define BS_CAP           (CAPACITY+1)

//
// auxiliary functions
//
inline int max( int a, int b ) { return a > b ? a : b; }
inline int min( int a, int b ) { return a < b ? a : b; }
double read_timer( )
{
    static int initialized = 0;
    static struct timeval start;
    struct timeval end;
    if( !initialized )
    {
        gettimeofday( &start, NULL );
        initialized = 1;
    }
    gettimeofday( &end, NULL );
    return (end.tv_sec - start.tv_sec) + 1.0e-6 * (end.tv_usec - start.tv_usec);
}

//
//  solvers
//

int build_table_serial( int nitems, int cap, shared [] int *T, int *w, int *v )
{
    int wj = w[0], vj = v[0];
    for( int i = 0;  i <  wj;  i++ ) T[i] = 0;
    for( int i = wj; i <= cap; i++ ) T[i] = vj;
    
    for( int j = 1; j < nitems; j++ ) 
    {
        int wj = w[j], vj = v[j];
        for( int i = 0;  i <  wj;  i++ ) T[i+cap+1] = T[i];
        for( int i = wj; i <= cap; i++ ) T[i+cap+1] = max( T[i], T[i-wj] + vj );
        
        T += cap+1;
    }
    
    return T[cap];
}

// Single thread, calculate the best value given the best values for a splitting of the items into two sets
void merge_rows( int cap, int *R1, int *R2, shared [BS_CAP] int *T, shared [BS_CAP] int *splits, int last ) {
  if( last ) {
    // For the last merge we only need to calculate at full capacity
    int i = cap;
    int best = -1;
    int split = -1;
    for( int w1 = 0; w1 <= i; w1++ ) {
      int option = R1[w1] + R2[i-w1];
      if( option > best ) {
        best = option;
	split = w1;
      }
    }
    T[i] = best;
    splits[i] = split;
    return;
  }
  for( int i = 0; i <= cap; i++ ) {
    int best = -1;
    int split = -1;
    for( int w1 = 0; w1 <= i; w1++ ) {
      int option = R1[w1] + R2[i-w1];
      if( option > best ) {
        best = option;
	split = w1;
      }
    }
    T[i] = best;
    splits[i] = split;
  }
}

// This code only works if THREADS is a power of 2
int build_table_parallel( int nitems, int cap, shared [] int * shared [1] T[], shared [BS_CAP] int *TT, shared [BS_CAP] int *splits, shared [BS_ITEMS] int *w, shared [BS_ITEMS] int *v ) {
  // First have each process build its table and put that in T

  int items_this_thread = (nitems + THREADS - 1) / THREADS;
  int first_item = items_this_thread * MYTHREAD;
  if (MYTHREAD == THREADS-1)
    items_this_thread = nitems - (THREADS-1)*items_this_thread;
  // This arithmetic is not right for small numbers.  This will catch errors until we fix it
  assert(items_this_thread > 0);
  int items_next_thread = items_this_thread;
  if( MYTHREAD == THREADS-2)
    items_next_thread = nitems - (THREADS-1)*items_next_thread;  
  int tmp = build_table_serial( items_this_thread, cap, T[MYTHREAD], (int*)(w+first_item), (int*)(v+first_item) );
  //printf("%d %d nitems %d\n", MYTHREAD, tmp, items_this_thread);
  upc_barrier;

  // Now do a binary merge on them into TT.
  if( THREADS == 1 ) {
    merge_rows( cap,  (int *)(T[MYTHREAD]+(cap+1)*(items_this_thread-1)),  (int*)(T[MYTHREAD+1]+(cap+1)*(items_next_thread-1)), TT+MYTHREAD/2*(cap+1), splits+MYTHREAD/2*(cap+1), 0 );
    return T[MYTHREAD][(cap+1)*(items_this_thread)-1];
  }

  int *row1 = (int *)malloc((cap+1)*sizeof(int));
  int *row2 = (int *)malloc((cap+1)*sizeof(int));
  int *arg1, *arg2;
  // Do the first step separately
  if( MYTHREAD % 2 ==  0 ) {
    upc_memget(row2, T[MYTHREAD+1]+(cap+1)*(items_next_thread-1), (cap+1)*sizeof(int));
    merge_rows( cap, (int*)(T[MYTHREAD]+(cap+1)*(items_this_thread-1)), row2, TT+MYTHREAD/2*(cap+1), splits+MYTHREAD/2*(cap+1), THREADS==2 );
}
  upc_barrier;

  // Then the rest

  for( int n = THREADS/2; n > 1; n/=2 ) {
    if( MYTHREAD % (2*THREADS/n) == 0 ) {
      int rank = MYTHREAD/(2*THREADS/n);
      //printf("n=%d, rank=%d, Merging from %d and %d to %d\n", n, rank, 2*rank*(cap+1), (2*rank+1)*(cap+1), (rank+n)*(cap+1));
      if( upc_threadof(TT+2*rank*(cap+1)) == MYTHREAD )
        arg1 = (int *)(TT+2*rank*(cap+1));
      else {
        upc_memget(row1, TT+2*rank*(cap+1), (cap+1)*sizeof(int));
        arg1 = row1;
      }
      if( upc_threadof(TT+(2*rank+1)*(cap+1)) == MYTHREAD )
        arg2 = (int *)(TT+2*rank*(cap+1));
      else {
        upc_memget(row2, TT+(2*rank+1)*(cap+1), (cap+1)*sizeof(int));
        arg2 = row2;
      }
      merge_rows( cap, arg1, arg2, TT+(rank+n)*(cap+1), splits+(rank+n)*(cap+1) , n==2);
    }
    upc_barrier;

    TT += (cap+1)*n;
    splits += (cap+1)*n;

  }

  return TT[cap];
}

//scratch should be a cyclic array of size THREADS.
void backtrack_parallel( int nitems, int cap, shared [] int * shared [1] T[], shared [BS_CAP] int *splits, shared [BS_ITEMS] int *w, shared [BS_ITEMS] int *u, shared [1] int *scratch ) {
  if( MYTHREAD == 0 )
    scratch[0] = cap;
  upc_barrier;

  splits += (cap+1)*(THREADS-2);
  for( int n = 2; n <= THREADS; n*=2 ) {
    if( MYTHREAD % (2*THREADS/n) == 0 ) {
     int rank = MYTHREAD/(2*THREADS/n);
     int t = scratch[MYTHREAD];
     int s = splits[rank*(cap+1)+t];
     scratch[MYTHREAD] = s;
     scratch[MYTHREAD+THREADS/n] = t-s;
     //printf("%d Splitting %d into %d and %d\n", MYTHREAD, t, s, t-s);
    }
    splits -= (cap+1)*n;
    upc_barrier;
  }
  // Now scratch should hold the capacity that will be devoted to each processors items.

  // This better be the same as in build_table_parallel.  Really it should be moved to a function
  int items_this_thread = (nitems + THREADS - 1) / THREADS;
  int first_item = items_this_thread * MYTHREAD;
  if (MYTHREAD == THREADS-1)
    items_this_thread = nitems - (THREADS-1)*items_this_thread;
  int i = items_this_thread * (cap+1) - 1 - cap + scratch[MYTHREAD];
  u += first_item;
  w += first_item;
  for( int j = items_this_thread-1; j > 0; j-- ) {
    u[j] = T[MYTHREAD][i] != T[MYTHREAD][i-cap-1];
    i -= cap+1 + (u[j] ? w[j] : 0);
  }
  u[0] = T[MYTHREAD][i] != 0;
  upc_barrier;
}

void backtrack( int nitems, int cap, shared int *T, shared int *w, shared int *u )
{
    int i, j;
    
    if( MYTHREAD != 0 )
        return;
    
    i = nitems*(cap+1) - 1;
    for( j = nitems-1; j > 0; j-- )
    {
        u[j] = T[i] != T[i-cap-1];
        i -= cap+1 + (u[j] ? w[j] : 0 );
    }
    u[0] = T[i] != 0;
}

//
//  serial solver to check correctness
//
int solve_serial( int nitems, int cap, shared [BS_ITEMS] int *w, shared [BS_ITEMS] int *v )
{
    int i, j, best, *allocated, *T, wj, vj;
    
    //alloc local resources
    T = allocated = malloc( nitems*(cap+1)*sizeof(int) );
    if( !allocated )
    {
        fprintf( stderr, "Failed to allocate memory" );
        upc_global_exit( -1 );
    }
    
    //build_table locally
    wj = w[0];
    vj = v[0];
    for( i = 0;  i <  wj;  i++ ) T[i] = 0;
    for( i = wj; i <= cap; i++ ) T[i] = vj;
    for( j = 1; j < nitems; j++ ) 
    {
        wj = w[j];
        vj = v[j];
        for( i = 0;  i <  wj;  i++ ) T[i+cap+1] = T[i];
        for( i = wj; i <= cap; i++ ) T[i+cap+1] = max( T[i], T[i-wj]+vj );
        T += cap+1;
    }
    best = T[cap];
    
    //free resources
    free( allocated );
    
    return best;
}

//
//  benchmarking program
//

shared [] int * shared [1] total_int[THREADS];

int main( int argc, char** argv )
{
    assert(THREADS == NUM_THREADS);

    int i, best_value, best_value_serial, total_weight, nused, total_value;
    double seconds;
    shared [BS_ITEMS] int *weight;
    shared [BS_ITEMS] int *value;
    shared [BS_ITEMS] int *used;
    shared [BS_CAP] int *total;
    shared [BS_CAP] int *splits;
    shared [1] int *scratch;

    //these constants have little effect on runtime
    int max_value  = 1000;
    int max_weight = 1000;
    
    //these set the problem size
    //int capacity   = 999;
    //int nitems     = 5000;
    int capacity = CAPACITY;
    int nitems = NITEMS;

    srand48( (unsigned int)time(NULL) + MYTHREAD );
    //allocate distributed arrays, use cyclic distribution
    weight = (shared [BS_ITEMS] int *) upc_all_alloc( THREADS, (nitems + THREADS - 1) / THREADS * sizeof(int) );
    value  = (shared [BS_ITEMS] int *) upc_all_alloc( THREADS, (nitems + THREADS - 1) / THREADS * sizeof(int) );
    used   = (shared [BS_ITEMS] int *) upc_all_alloc( THREADS, (nitems + THREADS - 1) / THREADS * sizeof(int) );
    // We should probably rethink the affinity of this / resize it.
    total  = (shared [BS_CAP] int *) upc_all_alloc( THREADS, (capacity+1) * sizeof(int) );
    total_int[MYTHREAD] = (shared [] int *)upc_alloc( (nitems+THREADS-1)/THREADS*(capacity+1)*sizeof(int) );
    splits  = (shared [BS_CAP] int *) upc_all_alloc( nitems, (capacity+1)*sizeof(int) );
    scratch = (shared [1] int *) upc_all_alloc( THREADS, sizeof(int) );
    if( !weight || !value || !total || !used )
    {
        fprintf( stderr, "Failed to allocate memory" );
        upc_global_exit( -1 );
    }
    
    // init
    max_weight = min( max_weight, capacity );//don't generate items that don't fit into bag
    upc_forall( i = 0; i < nitems; i++; &weight[i] )
    {
        weight[i] = 1 + (lrand48()%max_weight);
        value[i]  = 1 + (lrand48()%max_value);
    }
    upc_barrier;

    // time the solution
    seconds = read_timer( );
    
    best_value = build_table_parallel( nitems, capacity, total_int, total, splits, weight, value );
    backtrack_parallel( nitems, capacity, total_int, splits, weight, used, scratch );
    
    seconds = read_timer( ) - seconds;
    
    // check the result
    if( MYTHREAD == 0 )
    {
        printf( "%d items, capacity: %d, time: %g\n", nitems, capacity, seconds );
        
        best_value_serial = solve_serial( nitems, capacity, weight, value );
        
        total_weight = nused = total_value = 0;
        for( i = 0; i < nitems; i++ )
            if( used[i] )
            {
                nused++;
                total_weight += weight[i];
                total_value += value[i];
            }
        
        printf( "%d items used, value %d, weight %d\n", nused, total_value, total_weight );
        
        if( best_value != best_value_serial || best_value != total_value || total_weight > capacity )
            printf( "WRONG SOLUTION %d %d %d %d %d\n", best_value, best_value_serial, total_value, total_weight, capacity );
        
        //release resources
        upc_free( weight );
        upc_free( value );
        upc_free( total );
        upc_free( used );
	upc_free( splits );
	upc_free( scratch );
    }
    
    return 0;
}
